from functools import cached_property
import logging
from typing import ClassVar
from pathlib import Path
import polars as pl
from pysam import VariantFile  # pylint: disable=E0611
from pycircos import Garc, Gcircle
from pydantic import BaseModel, Field

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger("MTDD")


class Image(BaseModel):
    ref_size: ClassVar[int] = 16569
    depth_file: Path = Field(description="Path to depth file")
    genepred_file: Path = Field(description="Path to gene prediction file")
    vcf_file: Path = Field(description="Path to VCF file")
    image_file: Path = Field(description="Path to image file")
    circle: Gcircle = Field(default_factory=lambda: Gcircle(figsize=(8, 8)), description="Circle")

    class Config:
        arbitrary_types_allowed = True

    @cached_property
    def depth(self) -> pl.DataFrame:
        return pl.read_csv(self.depth_file, separator="\t")

    @cached_property
    def transcripts(self) -> pl.DataFrame:
        keys = ["name", "chrom", "gene", "strand", "start", "end", "unk"]
        records = (
            pl.read_csv(
                self.genepred_file,
                has_header=False,
                separator="\t",
                columns=[2, 3, 4, 5, 6, 7, 12],
                new_columns=["chrom", "strand", "start", "end", "cds_start", "cds_end", "gene"],
            )
            .with_columns(pl.when(pl.col("cds_end").eq(pl.col("cds_start"))).then(pl.lit("T")).otherwise(pl.lit("F")).alias("unk"))
            .sort(["chrom", "start"], descending=False)
            .with_row_index(name="name", offset=1)
            .with_columns(pl.format("gene{}", pl.col("name")).alias("name"))
        )
        gap_starts, gap_ends = records["end"][:-1].to_list(), records["start"][1:].to_list()
        if gap_starts[0] > 0:
            gap_ends = gap_starts[:1] + gap_ends
            gap_starts = [0] + gap_starts
        if gap_ends[-1] < self.ref_size:
            gap_starts += gap_ends[-1:]
            gap_ends += [self.ref_size]
        gap_records = (
            pl.DataFrame({"start": gap_starts, "end": gap_ends})
            .with_columns(
                pl.lit("chrM").alias("chrom"),
                pl.lit("").alias("strand"),
                pl.lit(" ").alias("gene"),
                pl.lit("").alias("unk"),
            )
            .filter(pl.col("start").lt(pl.col("end")))
            .with_row_index(name="name", offset=1)
            .with_columns(pl.format("gene{}", pl.col("name")).alias("name"))
        )
        return pl.concat([records.select(keys), gap_records.select(keys)]).sort(["chrom", "start"], descending=False)

    def plot_gene(self):
        index = 0
        for row in self.transcripts.iter_rows(named=True):
            match row["unk"]:
                case "T":
                    color = "#DF8D8F"
                case "F":
                    color = "#7EA2ED"
                case _:
                    color = "#E6E6E6"
            if row["end"] - row["start"] > 300:
                label_pos = 60
                index = 0
            else:
                if row["gene"].strip():
                    label_pos = 60 if index % 2 else -60
                    index += 1
                else:
                    label_pos = 60
            arc = Garc(
                arc_id=row["name"],
                size=row["end"] - row["start"],
                interspace=0,
                raxis_range=(935, 985),
                labelposition=label_pos,
                label_visible=True,
                label=row["gene"].replace("MT-", ""),
                facecolor=color,
            )
            if arc.size > 1:
                self.circle.add_garc(arc)
        self.circle.set_garcs()

    def plot_depth(self):
        for row in self.transcripts.iter_rows(named=True):
            t_records = self.depth.filter(pl.col("pos").is_between(row["start"], row["end"], closed="right"))
            self.circle.fillplot(garc_id=row["name"], data=t_records["depth"], facecolor="r", raxis_range=(785, 835))

    def plot_cnv(self):
        with VariantFile(self.vcf_file) as reader:
            for record in reader:
                start_row = self.transcripts.filter(pl.col("start").lt(record.start) & pl.col("end").ge(record.start)).row(0, named=True)
                end_row = self.transcripts.filter(pl.col("start").lt(record.stop) & pl.col("end").ge(record.stop)).row(0, named=True)
                start, end = record.start - start_row["start"], record.stop - end_row["start"]
                su = record.samples[0]["SU"]
                # print(str(record), start_row, end_row, start, end, su)
                self.circle.chord_plot(
                    [start_row["name"], start, start + su, 735],
                    [end_row["name"], end, end + su, 735],
                    facecolor="#ff8c0080",
                )

    def run(self):
        logging.info("Do Image to %s", self.image_file)
        self.plot_gene()
        self.plot_depth()
        self.plot_cnv()
        self.circle.save(file_name=str(self.image_file.with_suffix("")), format=self.image_file.suffix.lstrip("."))
